package com.cjc.common.rtsp;

import ai.djl.ModelException;
import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.TranslateException;
import com.cjc.common.rtsp.utils.FaceDetection;
import com.cjc.common.rtsp.utils.FaceMaskDetect;
import com.cjc.common.rtsp.utils.OpenCVImageUtil;
import org.bytedeco.javacv.FFmpegFrameGrabber;
import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.bytedeco.opencv.global.opencv_imgcodecs;
import org.bytedeco.opencv.opencv_core.Mat;
import org.bytedeco.opencv.opencv_core.Point;
import org.bytedeco.opencv.opencv_core.Rect;
import org.bytedeco.opencv.opencv_core.Scalar;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import java.awt.image.BufferedImage;
import java.io.IOException;
import java.util.List;

import static org.bytedeco.opencv.global.opencv_imgproc.*;

/**
 * 定时任务
 *
 * @Author: 陈江灿
 * @CreateTime: 2025-01-20
 */
@Component
public class RtspManager {

    @Value("${app.save-path}")
    private String savePath;

    @Scheduled(fixedRate = 10000)
    public void processAllRtspStreams() throws Exception {
//        String rtsp ="rtsp://admin:hx147258@192.168.158.65:554/Streaming/Channels/2";
//        faceMaskDetection(rtsp);
    }

    public void faceMaskDetection(String input)
            throws IOException, ModelException, TranslateException {
        float shrink = 0.5f;
        float threshold = 0.7f;
        Criteria<Image, DetectedObjects> criteria = new FaceDetection().criteria(shrink, threshold);
        Criteria<Image, Classifications> maskCriteria = new FaceMaskDetect().criteria();
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(input);
        if (input.indexOf("rtsp") > -1) {
            grabber.setFormat("rtsp");
            grabber.setOption("allowed_media_types", "video");
            grabber.setOption("rtsp_transport", "tcp");
            grabber.setOption("rtsp_flags", "prefer_tcp");
            grabber.setOption("stimeout", "3000000");
        }
        grabber.start();
        OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat();
        try (ZooModel model = ModelZoo.loadModel(criteria);
             Predictor<Image, DetectedObjects> predictor = model.newPredictor();
             ZooModel classifyModel = ModelZoo.loadModel(maskCriteria);
             Predictor<Image, Classifications> classifier = classifyModel.newPredictor()) {
            Frame frame = grabber.grabImage();
            if (frame != null) {
                Mat img = converter.convert(frame);
                BufferedImage buffImg = OpenCVImageUtil.mat2BufferedImage(img);
                Image image = ImageFactory.getInstance().fromImage(buffImg);
                int imageWidth = image.getWidth();
                int imageHeight = image.getHeight();
                DetectedObjects detections = predictor.predict(image);
                List<DetectedObjects.DetectedObject> items = detections.items();
                boolean hasUnmaskedFace = false;
                for (DetectedObjects.DetectedObject item : items) {
                    Image subImg = getSubImage(image, item.getBoundingBox());
                    Classifications classifications = classifier.predict(subImg);
                    String className = classifications.best().getClassName();
                    if ("NO MASK".equals(className)) {
                        hasUnmaskedFace = true;
                    }
                    BoundingBox box = item.getBoundingBox();
                    Rectangle rectangle = box.getBounds();
                    int x = (int) (rectangle.getX() * imageWidth);
                    int y = (int) (rectangle.getY() * imageHeight);
                    Rect face =
                            new Rect(
                                    x,
                                    y,
                                    (int) (rectangle.getWidth() * imageWidth),
                                    (int) (rectangle.getHeight() * imageHeight));
                    rectangle(img, face, new Scalar(0, 0, 255, 1));
                    int pos_x = Math.max(face.tl().x() - 10, 0);
                    int pos_y = Math.max(face.tl().y() - 10, 0);
                    putText(
                            img,
                            className,
                            new Point(pos_x, pos_y),
                            FONT_HERSHEY_COMPLEX,
                            1.0,
                            new Scalar(0, 0, 255, 2.0));
                }
                // 后续改为 上传图片即可
                if (hasUnmaskedFace) {
                    String outputPath = savePath + "frame_" + System.currentTimeMillis() + ".jpg";
                    opencv_imgcodecs.imwrite(outputPath, img);
                    System.out.println("检测到未戴口罩，保存图像: " + outputPath);
                }
            }
        } finally {
            grabber.close();
        }
    }

    private static int[] extendSquare(
            double xmin, double ymin, double width, double height, double percentage) {
        double centerx = xmin + width / 2;
        double centery = ymin + height / 2;
        double maxDist = Math.max(width / 2, height / 2) * (1 + percentage);
        return new int[] {(int) (centerx - maxDist), (int) (centery - maxDist), (int) (2 * maxDist)};
    }

    private static Image getSubImage(Image img, BoundingBox box) {
        Rectangle rect = box.getBounds();
        int width = img.getWidth();
        int height = img.getHeight();
        int[] squareBox =
                extendSquare(
                        rect.getX() * width,
                        rect.getY() * height,
                        rect.getWidth() * width,
                        rect.getHeight() * height,
                        0); // 0.18

        if (squareBox[0] < 0) squareBox[0] = 0;
        if (squareBox[1] < 0) squareBox[1] = 0;
        if (squareBox[0] > width) squareBox[0] = width;
        if (squareBox[1] > height) squareBox[1] = height;
        if ((squareBox[0] + squareBox[2]) > width) squareBox[2] = width - squareBox[0];
        if ((squareBox[1] + squareBox[2]) > height) squareBox[2] = height - squareBox[1];
        return img.getSubImage(squareBox[0], squareBox[1], squareBox[2], squareBox[2]);
    }
}
